import numpy as np

import matplotlib.pyplot as plt
import operator


# 给出训练数据以及对应的类别
def createDataSet():
    group = np.array([[1.0, 2.0], [1.2, 0.1], [0.1, 1.4], [0.3, 3.5], [1.1, 1.0], [0.5, 1.5]])
    labels = np.array(['A', 'A', 'B', 'B', 'A', 'B'])
    return group, labels


def kNN_classify(k, dis, X_train, x_train, Y_test):
    assert dis == 'E' or dis == 'M', 'dis must E or M，E代表欧拉距离，M代表曼哈顿距离'
    num_test = Y_test.shape[0]  # 测试样本的数量
    labellist = []
    '''
    使用欧拉公式作为距离度量
    '''
    if dis == 'E':
        for i in range(num_test):
            # 实现欧拉距离公式
            distances = np.sqrt(np.sum(((X_train - np.tile(Y_test[i], (X_train.shape[0], 1))) ** 2), axis=1))
            nearest_k = np.argsort(distances)  # 距离由小到大进行排序，并返回index值
            topK = nearest_k[:k]  # 选取前k个距离
            classCount = {}
            for i in topK:  # 统计每个类别的个数
                classCount[x_train[i]] = classCount.get(x_train[i], 0) + 1
            sortedClassCount = sorted(classCount.items(), key=operator.itemgetter(1), reverse=True)
            labellist.append(sortedClassCount[0][0])
        return np.array(labellist)


# 使用曼哈顿公式作为距离度量
# 读者自行补充完成


def mniset_download():
    import torch
    from torch.utils.data import DataLoader
    import torchvision.datasets as dsets
    import torchvision.transforms as transforms
    batch_size = 100
    # MNIST dataset
    train_dataset = dsets.MNIST(root='./pymnist',  # 选择数据的根目录
                                train=True,  # 选择训练集
                                transform=None,  # 不考虑使用任何数据预处理
                                download=True)  # 从网络上download图片
    test_dataset = dsets.MNIST(root='./pymnist',  # 选择数据的根目录
                               train=False,  # 选择测试集
                               transform=None,  # 不考虑使用任何数据预处理
                               download=True)  # 从网络上download图片
    # 加载数据
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                               batch_size=batch_size,
                                               shuffle=True)  # 将数据打乱
    test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                              batch_size=batch_size,
                                              shuffle=True)
    print("train_data", train_dataset.train_data.size())
    print("train_labels", train_dataset.train_labels.size())
    return train_loader, test_loader

    # import matplotlib.pyplot as plt
    # digit = train_loader.dataset.train_data[0]
    # plt.imshow(digit, cmap=plt.cm.binary)
    # plt.show()
    # print(train_loader.dataset.train_labels[0])


def dataset_test():
    train_loader, test_loader = mniset_download()
    train_loader.dataset.train_data.numpy()
    train_loader.dataset.train_data.numpy()


if __name__ == '__main__':
    # group,labels = createDataSet()
    # y_p = kNN_classify(1,'E',group,labels,np.array([[1.0,2.1],[0.4,2.0]]))
    # print(y_p)
    dataset_test()